package org.jie.guider.framework.springboot;

import java.util.concurrent.BlockingQueue;
import java.util.concurrent.RejectedExecutionException;
import java.util.concurrent.RejectedExecutionHandler;
import java.util.concurrent.SynchronousQueue;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * @author LiJie
 */
public class BlockThreadPool {

    private String threadGroupName;
    public ThreadPoolExecutor pool;
    private int corePoolSize;
    private int maxPoolSize;
    private long keepAliveTime;
    private TimeUnit unit;
    private BlockingQueue<Runnable> workQueue = new SynchronousQueue<>();
    /**
     * 线程池饱和时入队阻塞时间
     */
    private long waitTime;
    /**
     * 线程池饱和时入队阻塞时间单位
     */
    private TimeUnit waitUnit;

    /**
     * @param poolSize        定长线程
     * @param threadGroupName 线程组名称
     */
    public BlockThreadPool(int poolSize, String threadGroupName,
                           long waitTime, TimeUnit waitUnit) {
        this(poolSize, poolSize, 30, TimeUnit.SECONDS, threadGroupName, waitTime, waitUnit);
    }

    public BlockThreadPool(int corePoolSize, int maxPoolSize,
                           long keepAliveTime, TimeUnit unit,
                           String threadGroupName,
                           long waitTime, TimeUnit waitUnit) {

        this.corePoolSize = corePoolSize;
        this.maxPoolSize = maxPoolSize;
        this.keepAliveTime = keepAliveTime;
        this.unit = unit;
        this.waitTime = waitTime;
        this.waitUnit = waitUnit;
        this.threadGroupName = threadGroupName;
        init();
    }

    public void execute(Runnable task) {
        pool.execute(task);
    }

    public void shutdown() {
        pool.shutdown();
    }

    public boolean isShutdown() {
        return pool.isShutdown();
    }

    private void init() {
        ThreadFactory threadFactory = new ThreadFactoryInner(threadGroupName);
        RejectedExecutionHandler rejectedHandler =
            new RejectHandlerInner(waitTime, waitUnit);

        pool = new ThreadPoolExecutor
            (this.corePoolSize, this.maxPoolSize, this.keepAliveTime, this.unit,
             workQueue, threadFactory, rejectedHandler);
        pool.prestartAllCoreThreads();
    }

    private static class ThreadFactoryInner implements ThreadFactory {

        private final AtomicInteger poolNumber = new AtomicInteger(1);
        private ThreadGroup threadGroup;
        private String threadGroupName;

        public ThreadFactoryInner(String threadGroupName) {
            this.threadGroupName = threadGroupName;
            this.threadGroup = new ThreadGroup(threadGroupName);
        }

        @Override
        public Thread newThread(Runnable r) {
            String threadName = threadGroupName + "-" + poolNumber.getAndIncrement();
            return new Thread(threadGroup, r, threadName);
        }
    }

    private static class RejectHandlerInner implements RejectedExecutionHandler {
        /**
         * 线程池饱和时入队阻塞时间
         */
        private long waitTime;
        /**
         * 线程池饱和时入队阻塞时间单位
         */
        private TimeUnit waitUnit;

        public RejectHandlerInner(long waitTime, TimeUnit waitUnit) {
            this.waitTime = waitTime;
            this.waitUnit = waitUnit;
        }

        @Override
        public void rejectedExecution(Runnable r, ThreadPoolExecutor executor) {
            if (executor.isShutdown()) {
                throw throwException(r, executor, "pool shutdown", null);
            }
            try {
                boolean f = executor.getQueue().offer(r, waitTime, waitUnit);
                if (!f) {
                    throw throwException(r, executor, "enqueue timeout", null);
                }
            } catch (InterruptedException ex) {
                throw throwException(r, executor, "enqueue exception", ex);
            }
        }

        private RejectedExecutionException throwException
            (Runnable r, ThreadPoolExecutor executor, String exMsg, Exception e) {
            if (e != null) {
                return new RejectedExecutionException
                    ("Task " + r.toString() + " rejected[" + exMsg + "] from "
                         + executor.toString() + " : ", e);
            } else {
                return new RejectedExecutionException
                    ("Task " + r.toString() + " rejected[" + exMsg + "] from "
                         + executor.toString());
            }
        }
    }
}
